# -*- coding: utf-8 -*-
"""
@Time    : 2024/10/22 11:26 
@Author  : ZhangShenao 
@File    : generator.py 
@Desc    : 生成器

将用户的原始提问Query,与top-k个语义相似的文档块进行整合,构建完整的Prompt
并利用LLM的生成能力,生成最终回答
"""
from langchain_core.documents import Document
from langchain_core.language_models import BaseChatModel
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate


def generating(query: str, docs: list[Document], llm: BaseChatModel) -> None:
    """
    Generating 生成过程
    :param query: 原始提问
    :param docs: 相关文档内容
    :param llm: 大语言模型
    :return: 最终生成结果
    """

    print("Generating开始执行!")

    # 构造相关文档上下文
    docs = [doc.page_content for doc in docs]
    context = "\n\n".join(docs)

    # 构造Prompt
    human_msg = f"""用户原始提问: {query}
    以下是相关参考文档: 
    {context}
    
    """
    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是一个功能强大的AI助手,你的任务是回答用户的问题"),
        ("human", human_msg)
    ])

    # 构造Chain
    chain = prompt | llm | StrOutputParser()

    # 流式调用Chain,返回最终生成结果
    reply = chain.stream({})
    print("Generating执行完成!")
    print("最终生成结果: ")

    for content in reply:
        print(content, end="")
    print("")
    print("Generating执行完成！")
